import numpy as np
import torch
import matplotlib.pyplot as plt

def trainer():
    # 准备数据
    x = torch.linspace(0, 100, 100).type(torch.FloatTensor)
    rand = torch.randn(100) * 10
    y = x + rand
    x_train = x[:-10]
    x_test = x[-10:]
    y_train = y[:-10]
    y_test = y[-10:]

    #开始训练
    a, b = training(x_train, y_train)
    #绘图
    draw(x_train, y_train, x_test, y_test, a, b)

# 梯度下降法
def training(x_train, y_train):
    # 随机a,b
    a = torch.rand(1, requires_grad=True)
    b = torch.rand(1, requires_grad=True)
    print('a:', a, 'b:', b)

    # 步长
    learning_rate = 0.0001

    for i in range(1000):
        # 拟合直线方程：y=ax+b
        predictions = a.expand_as(x_train) * x_train + b.expand_as(x_train)
        # 损失函数
        loss = torch.mean((y_train - predictions) ** 2)
        # print("loss: ", loss)
        loss.backward()
        a.data.add_(-learning_rate * a.grad.data)
        b.data.add_(-learning_rate * b.grad.data)
        print('a.grad:', a.grad.data, 'b.grad:', b.grad.data, 'a.data', a.data, 'b.data', b.data)
        a.grad.data.zero_()
        b.grad.data.zero_()
    print('a:', a, 'b:', b)
    return a, b

def draw(x_train, y_train, x_test, y_test, a, b):
    x_data = x_train.data.numpy()
    x_pred = x_test.data.numpy()
    plt.figure(figsize=(10, 8))
    # 训练数据点
    xplot, = plt.plot(x_data, y_train.data.numpy(), 'o')
    # 测试数据点
    plt.plot(x_pred, y_test.data.numpy(), 'o')
    # 拟合数据
    x_data = np.r_[x_data, x_pred]
    yplot, = plt.plot(x_data, a.data.numpy() * x_data + b.data.numpy())
    # 预测数据
    plt.plot(x_pred, a.data.numpy() * x_pred + b.data.numpy(), 's')
    plt.xlabel('X')
    plt.ylabel('Y')
    str1 = str(a.data.numpy()[0]) + 'x ' + str(b.data.numpy()[0])
    plt.legend([xplot, yplot], ['Data', str1])
    plt.show()

trainer()